import random
from typing import Tuple, List

class PromptConstructor:
    def __init__(self, all_cn_diseases: List[str], stage_map: dict, stage_group: dict):
        self.all_cn_diseases = all_cn_diseases
        self.stage_map = stage_map
        self.stage_group = stage_group

    def construct_prompt(self, true_diag: str) -> Tuple[str, str]:
        rnd = random.random()

        if rnd < 0.3:
            return self._judge_prompt(true_diag)
        elif rnd < 0.8:
            return self._diagnosis_prompt(true_diag)
        else:
            return self._compare_prompt(true_diag)

    def _judge_prompt(self, true_diag: str) -> Tuple[str, str]:
        templates = [
            "图像中显示的是{}吗？请严格使用“是”或“否”回答。",
            "患者是否患有{}？请直接回答是或否。",
            "该眼底照片能诊断为{}吗？请明确回答是或否。",
            "请判断图像所示是否为{}，回答是或否。",
            "是否可以从该图像诊断出{}？请明确回答是或否。"
        ]

        is_positive = random.random() < 0.5
        diag = true_diag if is_positive else random.choice(
            [d for d in self.all_cn_diseases if d != true_diag])

        question = random.choice(templates).format(diag)
        answer = "是。" if is_positive else "否。"
        return question, answer

    def _diagnosis_prompt(self, true_diag: str) -> Tuple[str, str]:
        templates = [
            "请根据图像直接给出诊断结论。",
            "该患者的眼底情况如何？请直接诊断。",
            "根据图像，你认为患者的诊断结果是什么？",
            "请直接指出本图像的疾病诊断结果。",
            "分析眼底照片后，给出你的诊断结论。"
        ]

        question = random.choice(templates)
        answer_templates = [
            "本张图片的诊断结果是{}。",
            "该病例诊断为{}。",
            "经判断，本例患者为{}。",
            "此图像显示患者诊断为{}。",
            "眼底图片显示的疾病诊断为{}。"
        ]
        answer = random.choice(answer_templates).format(true_diag)
        return question, answer

    def _compare_prompt(self, true_diag: str) -> Tuple[str, str]:
        templates = [
            "图像更可能属于以下哪种疾病：{}或{}？请直接给出疾病名称。",
            "患者的眼底情况更符合{}还是{}？请直接回答疾病名称。",
            "根据图像，{}和{}中，哪种疾病的可能性更高？请直接回答疾病名称。",
            "从图像来看，患者的诊断更倾向于{}还是{}？直接给出诊断名称即可。",
            "请判断该图像更接近以下哪种疾病：{}或者{}？直接回答疾病名称即可。"
        ]

        wrong_candidates = [d for d in self.all_cn_diseases if d != true_diag]
        wrong_diag = random.choice(wrong_candidates)

        # 随机调换两个疾病顺序，避免顺序影响
        if random.random() < 0.5:
            diag_a, diag_b = true_diag, wrong_diag
        else:
            diag_a, diag_b = wrong_diag, true_diag

        question = random.choice(templates).format(diag_a, diag_b)
        answer = true_diag
        return question, answer

# 使用示例：
if __name__ == "__main__":
    all_cn_diseases = ["正常", "糖尿病视网膜病变早期", "糖尿病视网膜病变晚期", "青光眼", "视网膜裂孔", "老年黄斑变性"]
    stage_map = {
        "糖尿病视网膜病变早期": "糖尿病视网膜病变",
        "糖尿病视网膜病变晚期": "糖尿病视网膜病变"
    }
    stage_group = {
        "糖尿病视网膜病变": ["糖尿病视网膜病变早期", "糖尿病视网膜病变晚期"]
    }

    constructor = PromptConstructor(all_cn_diseases, stage_map, stage_group)
    for _ in range(10):
        q, a = constructor.construct_prompt("糖尿病视网膜病变晚期")
        print(f"Q: {q}\nA: {a}\n")
